import model
import pandas as pd
import pickle
from sklearn.model_selection import cross_val_predict,KFold,cross_validate
from sklearn.metrics import mean_absolute_error,r2_score
from sklearn.utils import shuffle

from sklearn.ensemble import RandomForestRegressor
from lightgbm import LGBMRegressor

def run_training(df, category_fields, numeric_fields, yfield, rs, regressor_name, **kwargs):
    x_part= df[category_fields+ numeric_fields]
    data_train, onehot_encoder = model.convert_traindf_to_matrix(x_part.loc[ :,   category_fields +  numeric_fields] ,
                                        category_fields, numeric_fields)
    target_train = df.loc[:, yfield].values

    data_train, target_train = shuffle(data_train, target_train, random_state=rs)

    regressor = regressor_name(**kwargs)
    regressor.fit(data_train, target_train)

    return regressor, onehot_encoder


if __name__ == "__main__":
    df = pd.read_csv("train_out.csv")
    from process import columns3
    df.columns = columns3

    category_fields = ["anchor", "dst_anchor", "carrierName"]
    numeric_fields=["speed1","relative_direction", "relative_distance", "relative_x", "relative_y",
     "abs_x", "abs_y", "sum_distance",  ]

    lgbm_sklearn = {
        'learning_rate': 0.1,
        'max_bin': 150,
        'num_leaves': 32,
        'max_depth': 11,

        'reg_alpha': 0.1,
        'reg_lambda': 0.2,

        'objective': 'regression',
        'n_estimators': 300,
    }

    regressor, onehot_encoder = run_training(df,category_fields, numeric_fields,  "remain_seconds",
                                             10,
                                             LGBMRegressor, **lgbm_sklearn )
    with open('pickle.regressor', 'wb') as f:
        pickle.dump(regressor, f)
    with open('pickle.encoder', 'wb') as f2:
        pickle.dump(onehot_encoder, f2)